import sqlite3
import pandas as pd
from FinSQL_01_generate import generate_sql
from FinSQL_02_query import query_db
from FinSQL_03_answer_from_SQL import generate_answer
from utils.configFinRAG import sql_examples_path
from utils.instances import TOKENIZER, LLM

# 初始化全局变量列表
example_lists = {
    'questions': [],
    'sqls': [],
    'infos': [],
    'fa': [],
    'tokens': []
}


def load_example_data():
    """加载示例数据"""
    examples = pd.read_csv(sql_examples_path, delimiter=",", header=0)
    for _, row in examples.iterrows():
        example_lists['questions'].append(row['问题'])
        example_lists['sqls'].append(row['SQL'])
        example_lists['infos'].append(row['资料'])
        example_lists['fa'].append(row['FA'])
        tokens = TOKENIZER(row['问题'])['input_ids']
        example_lists['tokens'].append(tokens)


def sql_retrieve_chain(query):
    """处理查询并生成回答"""
    if not example_lists['questions']:
        load_example_data()

    # 生成SQL语句
    result_prompt, sql = generate_sql(query, LLM, example_lists['questions'], example_lists['sqls'],
                                      example_lists['tokens'])

    # 连接数据库并执行查询
    with sqlite3.connect('/path/to/your/database.db') as conn:
        cs = conn.cursor()
        success_flag, exc_result = query_db(sql, cs)

    # 根据查询结果生成答案
    answer = generate_answer(query, exc_result, LLM, example_lists['questions'], example_lists['infos'],
                             example_lists['fa'], example_lists['tokens'])

    return answer